#include "PointCloudGenerator.h"

PointCloudGenerator::PointCloudGenerator(const std::string& calib_file)
:calib_file_(calib_file)
{
    debug = false;
    filter = true;

    warning = false;
    warning_point_unm = 5;

    width_ = 1280;
    height_ = 960;

    // 设置旋转，过滤默认参数
    //float tx, ty, tz, yaw, pitch, roll, 
    //float min_x, max_x, min_y, max_y, min_z, max_z
    setparams(0.0, 0.0, 0.0, 0.00, 0.0, 0.0, 
              0.0, 3.0, -5.0, 5.0, -3.0, 3.0);

    // 初始化相机参数
    if (!loadCameraParams(calib_file_)) {
        throw std::runtime_error("Failed to load camera parameters");
    }

    precomputeUndistortionMap();

    precomputeDenominator();

    cloud.reset(new pcl::PointCloud<pcl::PointXYZ>);
    cloud_downsample.reset(new pcl::PointCloud<pcl::PointXYZ>);
    remove_nan_point.reset(new pcl::PointCloud<pcl::PointXYZ>);
    final_filtered_cloud_.reset(new pcl::PointCloud<pcl::PointXYZ>);

    voxelgrid.reset(new pcl::VoxelGrid<pcl::PointXYZ>());
    approx_voxelgrid.reset(new pcl::ApproximateVoxelGrid<pcl::PointXYZ>);

    middle_height_start_ = 0;
    middle_height_end_ = 960;
    middle_height_ = 960;

    downsample_resolution = 0.005;
    filter_kernel_size = 5;
}

bool PointCloudGenerator::loadCameraParams(const std::string& file_path) {
    std::ifstream file(file_path);
    if (!file) return false;

    // 读取内参矩阵
    mtx_ = cv::Mat::eye(3, 3, CV_64F);
    if (!(file >> mtx_.at<double>(0,0) >> mtx_.at<double>(0,1) >> mtx_.at<double>(0,2)
            >> mtx_.at<double>(1,0) >> mtx_.at<double>(1,1) >> mtx_.at<double>(1,2)
            >> mtx_.at<double>(2,0) >> mtx_.at<double>(2,1) >> mtx_.at<double>(2,2))) {
            std::cout << "Failed to read camera matrix from file" << std::endl;
        return false;
    }

    // 读取畸变系数
    dist_ = cv::Mat::zeros(5, 1, CV_64F);
    if (!(file >> dist_.at<double>(0) >> dist_.at<double>(1) >> dist_.at<double>(2)
            >> dist_.at<double>(3) >> dist_.at<double>(4))) {
            std::cout << "Failed to read distortion coefficients from file" << std::endl;
        return false;
    }

    // 打印内参矩阵和畸变系数
    if(debug)
    {
        std::cout << "\nmtx_:\n" << mtx_ << "\n";
        std::cout << "dist_:\n" << dist_.t() << "\n";
    }
    return true;
}

void PointCloudGenerator::generate(const cv::Mat &depth_16u) 
{
    if (depth_16u.empty()) return;

    cv::Mat median_filtered;
    cv::medianBlur(depth_16u, median_filtered, filter_kernel_size); 
    if (median_filtered.empty())
    {
        return;
    }
    createCloud(median_filtered, undistorted_grid_, cloud);
    if(debug)
    {
        std::cout << " row: "  << " point number 1: " << cloud->size() << std::endl;
    }
    cloud->is_dense = false;   

    approx_voxelgrid->setInputCloud(cloud);
    approx_voxelgrid->setLeafSize(downsample_resolution, downsample_resolution, downsample_resolution);
    approx_voxelgrid->filter(*cloud_downsample);
    if(debug)
    {
        std::cout << "after downsample 2 point number: " << cloud_downsample->size() << std::endl;
    }

    if(filter)
    {
        pcl::PointCloud<pcl::PointXYZ>::Ptr transformed_cloud_(new pcl::PointCloud<pcl::PointXYZ>());
        pcl::transformPointCloud(*cloud_downsample, *transformed_cloud_, transform);
    
        // 组合条件滤波（优化后的区域过滤）
        pcl::ConditionAnd<pcl::PointXYZ>::Ptr range_cond(new pcl::ConditionAnd<pcl::PointXYZ>);
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("x", pcl::ComparisonOps::GT, min_x_)));
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("x", pcl::ComparisonOps::LT, max_x_)));
    
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("y", pcl::ComparisonOps::GT, min_y_)));
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("y", pcl::ComparisonOps::LT, max_y_)));
    
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("z", pcl::ComparisonOps::GT, min_z_)));
        range_cond->addComparison(pcl::FieldComparison<pcl::PointXYZ>::ConstPtr(
            new pcl::FieldComparison<pcl::PointXYZ>("z", pcl::ComparisonOps::LT, max_z_)));
    
        pcl::ConditionalRemoval<pcl::PointXYZ> cond_filter;
        cond_filter.setCondition(range_cond);
        cond_filter.setInputCloud(transformed_cloud_);
        cond_filter.filter(*final_filtered_cloud_);

        if(final_filtered_cloud_->size() > warning_point_unm)
        {
            warning = true;
        }
        else
        {
            warning = false;
        }
    }

}

void PointCloudGenerator::precomputeUndistortionMap() {

    cv::Mat grid = createPixelGrid();

    cv::Mat continuous_grid = grid.clone();
    continuous_grid.convertTo(continuous_grid, CV_32F);

    cv::undistortPoints(continuous_grid, undistorted_grid_, mtx_, dist_);

    undistorted_grid_ = undistorted_grid_.reshape(2, height_);
}

cv::Mat PointCloudGenerator::createPixelGrid() {
    cv::Mat grid(1, height_ * width_, CV_32FC2);
    cv::Vec2f* ptr = grid.ptr<cv::Vec2f>();

    for (int i = 0; i < height_; ++i) {
        for (int j = 0; j < width_; ++j) {
            ptr[i * width_ + j] = cv::Vec2f(static_cast<float>(j), static_cast<float>(i));
        }
    }

    return grid;
}

void PointCloudGenerator::precomputeDenominator() {
    denominator_.resize(height_ * width_);
    const cv::Vec2f* uv = undistorted_grid_.ptr<cv::Vec2f>();
    for (int i = 0; i < height_; ++i) {
        for (int j = 0; j < width_; ++j) {
            const int idx = i * width_ + j;
            const float map_x = uv[idx][0];
            const float map_y = uv[idx][1];
            denominator_[idx] = sqrt(1.0f + map_x * map_x + map_y * map_y)*1000;
        }
    }
}

void PointCloudGenerator::createCloud(const cv::Mat& depth, const cv::Mat& undistorted, pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud) 
{
    const cv::Vec2f* uv = undistorted.ptr<cv::Vec2f>();
    pcl::PointXYZ* data_ptr = cloud->points.data(); // 获取内存块首地址

    for (int i = 0; i < middle_height_ ; i++ ) 
    {
        const ushort* depth_row = depth.ptr<ushort>(i); // 预先获取当前行的指针
        for (int j = 0; j < middle_width_; j++) 
        {
            const int idx = (i + middle_height_start_) * width_ + j + middle_width_start_;
            const int linear_idx = i * middle_width_ + j; // 内存连续索引
            const ushort z_raw_ushort = depth_row[j]; 

            const float z_raw = static_cast<float>(z_raw_ushort);

            if (z_raw_ushort <= 0.001f) {
                data_ptr[linear_idx] = pcl::PointXYZ(NAN, NAN, NAN); 
                continue;
            }

            const float map_x = uv[idx][0];
            const float map_y = uv[idx][1];

            const float z_corrected = z_raw / denominator_[idx];
   
            data_ptr[linear_idx].x = z_corrected;
            data_ptr[linear_idx].y = -map_x * z_corrected;
            data_ptr[linear_idx].z = -map_y * z_corrected;
        }
    }
}

void PointCloudGenerator::setparams(float tx, float ty, float tz, float yaw, float pitch, float roll, 
                                    float min_x, float max_x, float min_y, float max_y, float min_z, float max_z)
{
    tx_ = tx;
    ty_ = ty;
    tz_ = tz;
    yaw_ = yaw;
    pitch_ = pitch;
    roll_ = roll;
    transform = Eigen::Affine3f::Identity();
    transform.translation() << tx, ty, tz;
    Eigen::AngleAxisf rollAngle(roll, Eigen::Vector3f::UnitX());
    Eigen::AngleAxisf pitchAngle(pitch, Eigen::Vector3f::UnitY());
    Eigen::AngleAxisf yawAngle(yaw, Eigen::Vector3f::UnitZ());
    transform.linear() = (yawAngle * pitchAngle * rollAngle).matrix();

    min_x_ = min_x;
    max_x_ = max_x;
    min_y_ = min_y;
    max_y_ = max_y; 
    min_z_ = min_z; 
    max_z_ = max_z;
}


